//////////////////////////////////////////////////////////////////
// (c) copyright 2010- by jeongnim kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications 
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// supported by 
//   national center for supercomputing applications, uiuc
//////////////////////////////////////////////////////////////////
// -*- c++ -*-
#include <dns.hpp>

using namespace std;
int main(int argc, char** argv)
{
  typedef  double real_type;
  const int n_var=3;
  int niters=10;
  int n=12288;
  int npx=512;
  int npy=512;
  bool debug=false;
  double dt=0.001;

  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-d")<a.size()) //get the z
      n=atoi(argv[++ic]);
    if(a.find("-npx")<a.size()) //x-partition
      npx=atoi(argv[++ic]);
    if(a.find("-npy")<a.size()) //x-partition
      npy=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-debug")<a.size())  
      debug=true;
    ++ic;
  }

  
  //timestep
  double dt_2=dt/2.0;
  double dt_3=dt/3.0;
  double dt_6=dt/6.0;

  DNS::Problem<double> cubic(n,npx,npy);
  
  typedef DNS::Container<double,3> array_t;

  ///velocity
  array_t V(cubic), V_r(cubic);
  ///righthand side
  array_t rhs(cubic);
  ///intermediate values for V duirng the stages
  vector<array_t*> Q;
  for(int i=0; i<2; ++i) Q.push_back(new array_t(cubic));

  //grep the total size
  int ngtot=static_cast<int>(rhs.size());
  cout << " Starting DNS simulation niters="<<niters << " " << cubic.size() << " " << rhs.size() << endl;

  //initialize V
  V.randomize();
  //initialize intermediate variables with V
  //*Q[0]=V;
  //*Q[1]=V;

  V.ifft(V_r);

  Timer clock;
  double dt_rk=0.0, dt_vor=0.0;
  for(int i=0; i<niters; ++i)
  {
    {//stage 1: V[1] is in k 
      //rhs_g = k x Qhat
      cubic.vorticity(V.data(),rhs.data());
      rhs.ifft();

      //V_r = V_r x rhs_r
      cubic.outerproduct(rhs.data(),V_r.data());
      V_r.fft(rhs);

      //rhs += -|k|^2 V
      cubic.laplacians(V.data(),rhs.data());

      //apply spherical aliasing: NOT IMPLEMENTED YET
      cubic.aliasing(rhs.data());

      DNS::axpy(V.data(),rhs.data(),dt_6,ngtot);
      DNS::axpy(Q[1]->data(),Q[0]->data(),rhs.data(),dt_2,ngtot);

      Q[1]->ifft(V_r);
    }

    { //stage 2
      cubic.vorticity(Q[1]->data(),rhs.data());
      rhs.ifft();

      cubic.outerproduct(rhs.data(),V_r.data());
      V_r.fft(rhs);

      //rhs += -|k|^2 V
      cubic.laplacians(Q[1]->data(),rhs.data());

      DNS::axpy(V.data(),rhs.data(),dt_3,ngtot);
      DNS::axpy(Q[0]->data(),Q[1]->data(),rhs.data(),dt_2,ngtot);

      Q[0]->ifft(V_r);
    }

    {//stage 3
      cubic.vorticity(Q[0]->data(),rhs.data());
      rhs.ifft();

      //V_r = V_r x rhs_r
      cubic.outerproduct(rhs.data(),V_r.data());
      V_r.fft(rhs);

      //rhs += -|k|^2 V
      cubic.laplacians(Q[0]->data(),rhs.data());

      DNS::axpy(V.data(),rhs.data(),dt_3,ngtot);
      DNS::axpy(Q[1]->data(),Q[0]->data(),rhs.data(),dt,ngtot);

      Q[1]->ifft(V_r);
    }

    {//stage 4
      cubic.vorticity(Q[1]->data(),rhs.data());
      dt_vor+=clock.elapsed();

      //V_r = V_r x rhs_r
      cubic.outerproduct(rhs.data(),V_r.data());
      V_r.fft(rhs);

      //rhs += -|k|^2 V
      cubic.laplacians(Q[1]->data(),rhs.data());

      DNS::axpy(V.data(),rhs.data(),dt,ngtot);
      V.ifft(V_r);
    }
  }

  double elpased=clock.elapsed();
  double fac=1.0/static_cast<double>(niters);
  printf("#tag n  npx  npy  rk4  \n");
  printf("dns %u %d %d %12.4e \n", n, npx, npy, elpased*fac);

  for(int i=0; i<Q.size(); ++i) delete Q[i];
  return 0;
}

